'''
Created on 2012-4-23

@author: nzhao
'''
import numpy as np

from Spin import Spin
from MultiSpin import MultiSpin
from SpinCluster import SpinCluster

class SpinBath(MultiSpin):
    '''
    spin bath
    '''


    def __init__(self, bathFile, centerSpin=[Spin()] ):
        '''
        load a spin bath 
        '''
        MultiSpin.__init__(self, cSpin=centerSpin)
     
        self.bathFile=bathFile
        spinInfo = self.readBathFile(bathFile)
        for spin in spinInfo:
            self.add_spin(spin)
        self.connectionMap=[]
    
    def readBathFile(self, filename):
        f=open(filename,'r')
        atomList=[line.split() for line in f][2:]
        
        spinInfo=[Spin(atom[0],np.array([float(atom[1]), float(atom[2]), float(atom[3])]) ) for atom in atomList]
        return spinInfo
    
    def clustering(self, maxSize=4, cutDistantce=6.0, cutStrength=0.1):
        self.calculate_interaction()
        self.calculate_coupling()
        
        self.connectionMap = [self.isConnect(i, j, cutStrength) \
                              for i in range(self.spinNum)\
                              for j in range(self.spinNum) if i > j]
        
        self.maxClusterSize = maxSize
        self.clusterList = [ [] for i in range(maxSize) ] 
        
        for spin in self.spinList:
            self.clusterList[0].append( SpinCluster( [spin] ) )
            
        for size in range(1,maxSize):
            for old_clst in self.clusterList[size - 1]:
                new_clst = self.clusterGrowFrom(old_clst)
                self.clusterList[size] = self.clusterList[size] + new_clst
        
            
    def clusterGrowFrom(self, old_clst):
        res = SpinCluster()
        return res

    def isConnect(self, i, j, cut):
        res = max( abs(self.interaction[self.index2(i, j)]) ) > cut
        return res

if __name__ == '__main__':
    cSpin=Spin()
    spinbath1=SpinBath(bathFile="/home/nzhao/research/computation/workspace/cce-engine/input/bath/C12Abund98_9/Abund11seed1.xyz")
    
    print spinbath1.spinNum
    print spinbath1.spinList[1].show()
    spinbath1.clustering()
    print spinbath1.interaction[:10]
    print spinbath1.couping[:5]
    print spinbath1.connectionMap[:3]